{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch import nn, optim, autograd\n",
    "from torch.nn import functional as F\n",
    "from torch.distributions.multivariate_normal import MultivariateNormal\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "#import input_data\n",
    "from sklearn.utils import shuffle as skshuffle\n",
    "from math import *\n",
    "from backpack import backpack, extend\n",
    "from backpack.extensions import KFAC, DiagHessian, DiagGGNMC\n",
    "from sklearn.metrics import roc_auc_score\n",
    "import scipy\n",
    "from tqdm import tqdm, trange\n",
    "from bpjacext import NetJac\n",
    "import pytest\n",
    "from DirLPA_utils import * \n",
    "import time\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "s = 123\n",
    "np.random.seed(s)\n",
    "torch.manual_seed(s)\n",
    "torch.cuda.manual_seed(s)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NOTE: DO NOT RUN THIS CODE: the function NetJac comes from a private repository that is not yet available for you. \n",
    "#If you search for it aggressively you might find out my identity. I would therefore prefer if you just looked at the provided results.\n",
    "\n",
    "#Also, Pytorch was updated and is now incompatible with NetJac. The results for the experiments however, can \n",
    "#still be found in the code. \n",
    "\n",
    "# since our package is not compatible with pytorch anymore the experiments with other integral approximations\n",
    "# have been conducted with a last-layer Laplace approximation of the network in the other jupyter notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device:  cuda\n",
      "cuda status:  True\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "cuda_status = torch.cuda.is_available()\n",
    "print(\"device: \", device)\n",
    "print(\"cuda status: \", cuda_status)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LPADirNN(num_classes=10, num_LL=256):\n",
    "    \n",
    "    features = torch.nn.Sequential(\n",
    "        torch.nn.Conv2d(1, 32, 5),\n",
    "        torch.nn.ReLU(),\n",
    "        torch.nn.MaxPool2d(2,2),\n",
    "        torch.nn.Conv2d(32, 64, 5),\n",
    "        torch.nn.ReLU(),\n",
    "        torch.nn.MaxPool2d(2,2),\n",
    "        torch.nn.Flatten(),\n",
    "        torch.nn.Linear(4 * 4 * 64, num_LL), #changed from 500\n",
    "        torch.nn.Linear(num_LL, num_classes)  #changed from 500\n",
    "    )\n",
    "    return(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE_TRAIN_MNIST = 32#64#128\n",
    "BATCH_SIZE_TEST_MNIST = 32#64#128\n",
    "MAX_ITER_MNIST = 6\n",
    "LR_TRAIN_MNIST = 10e-6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "MNIST_transform = torchvision.transforms.ToTensor()\n",
    "\n",
    "MNIST_train = torchvision.datasets.MNIST(\n",
    "        '~/data/mnist',\n",
    "        train=True,\n",
    "        download=True,\n",
    "        transform=MNIST_transform)\n",
    "\n",
    "mnist_train_loader = torch.utils.data.dataloader.DataLoader(\n",
    "    MNIST_train,\n",
    "    batch_size=BATCH_SIZE_TRAIN_MNIST,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    "\n",
    "MNIST_test = torchvision.datasets.MNIST(\n",
    "        '~/data/mnist',\n",
    "        train=False,\n",
    "        download=False,\n",
    "        transform=MNIST_transform)\n",
    "\n",
    "mnist_test_loader = torch.utils.data.dataloader.DataLoader(\n",
    "    MNIST_test,\n",
    "    batch_size=BATCH_SIZE_TEST_MNIST,\n",
    "    shuffle=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_model = LPADirNN(num_LL=256).cuda()\n",
    "loss_function = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "mnist_train_optimizer = torch.optim.Adam(mnist_model.parameters(), lr=1e-3, weight_decay=5e-4)\n",
    "MNIST_PATH = \"pretrained_weights/MNIST_pretrained_10_classes.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Training routine\n",
    "\n",
    "def train(model, train_loader, optimizer, max_iter, path, verbose=True):\n",
    "    max_len = len(train_loader)\n",
    "\n",
    "    for iter in range(max_iter):\n",
    "        for batch_idx, (x, y) in enumerate(train_loader):\n",
    "            \n",
    "            x, y = x.cuda(), y#.cuda()\n",
    "            \n",
    "            output = model(x)\n",
    "\n",
    "            accuracy = get_accuracy(output, y)\n",
    "\n",
    "            loss = loss_function(output, y)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            if verbose and batch_idx % 50 == 0:\n",
    "                print(\n",
    "                    \"Iteration {}; {}/{} \\t\".format(iter, batch_idx, max_len) +\n",
    "                    \"Minibatch Loss %.3f  \" % (loss) +\n",
    "                    \"Accuracy %.0f\" % (accuracy * 100) + \"%\"\n",
    "                )\n",
    "\n",
    "    print(\"saving model at: {}\".format(path))\n",
    "    torch.save(mnist_model.state_dict(), path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train(mnist_model, mnist_train_loader, mnist_train_optimizer, MAX_ITER_MNIST, MNIST_PATH, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading model from: pretrained_weights/MNIST_pretrained_10_classes.pth\n",
      "Batch 0/313 \tAccuracy 100%\n",
      "Batch 10/313 \tAccuracy 97%\n",
      "Batch 20/313 \tAccuracy 94%\n",
      "Batch 30/313 \tAccuracy 97%\n",
      "Batch 40/313 \tAccuracy 97%\n",
      "Batch 50/313 \tAccuracy 97%\n",
      "Batch 60/313 \tAccuracy 100%\n",
      "Batch 70/313 \tAccuracy 100%\n",
      "Batch 80/313 \tAccuracy 100%\n",
      "Batch 90/313 \tAccuracy 97%\n",
      "Batch 100/313 \tAccuracy 100%\n",
      "Batch 110/313 \tAccuracy 94%\n",
      "Batch 120/313 \tAccuracy 97%\n",
      "Batch 130/313 \tAccuracy 97%\n",
      "Batch 140/313 \tAccuracy 100%\n",
      "Batch 150/313 \tAccuracy 97%\n",
      "Batch 160/313 \tAccuracy 100%\n",
      "Batch 170/313 \tAccuracy 100%\n",
      "Batch 180/313 \tAccuracy 100%\n",
      "Batch 190/313 \tAccuracy 97%\n",
      "Batch 200/313 \tAccuracy 100%\n",
      "Batch 210/313 \tAccuracy 100%\n",
      "Batch 220/313 \tAccuracy 100%\n",
      "Batch 230/313 \tAccuracy 100%\n",
      "Batch 240/313 \tAccuracy 100%\n",
      "Batch 250/313 \tAccuracy 100%\n",
      "Batch 260/313 \tAccuracy 100%\n",
      "Batch 270/313 \tAccuracy 100%\n",
      "Batch 280/313 \tAccuracy 100%\n",
      "Batch 290/313 \tAccuracy 100%\n",
      "Batch 300/313 \tAccuracy 100%\n",
      "Batch 310/313 \tAccuracy 100%\n",
      "overall test accuracy on MNIST: 99.00 %\n"
     ]
    }
   ],
   "source": [
    "#predict in distribution\n",
    "MNIST_PATH = \"pretrained_weights/MNIST_pretrained_10_classes.pth\"\n",
    "\n",
    "#mnist_model = LPADirNN(x=100)\n",
    "mnist_model = LPADirNN(num_LL=256).cuda()\n",
    "print(\"loading model from: {}\".format(MNIST_PATH))\n",
    "mnist_model.load_state_dict(torch.load(MNIST_PATH))\n",
    "mnist_model.eval()\n",
    "\n",
    "acc = []\n",
    "\n",
    "max_len = len(mnist_test_loader)\n",
    "for batch_idx, (x, y) in enumerate(mnist_test_loader):\n",
    "\n",
    "    x, y = x.cuda(), y.cuda()\n",
    "    output = mnist_model(x)\n",
    "\n",
    "    accuracy = get_accuracy(output, y)\n",
    "    if batch_idx % 10 == 0:\n",
    "        print(\n",
    "            \"Batch {}/{} \\t\".format(batch_idx, max_len) + \n",
    "            \"Accuracy %.0f\" % (accuracy * 100) + \"%\"\n",
    "        )\n",
    "    acc.append(accuracy)\n",
    "\n",
    "avg_acc = np.mean(acc)\n",
    "print('overall test accuracy on MNIST: {:.02f} %'.format(avg_acc * 100))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "## play around with Backpack\n",
    "def get_Hessian_NN(model, train_loader, var0, device='cpu', verbose=True):\n",
    "    lossfunc = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    extend(lossfunc, debug=False)\n",
    "    extend(model, debug=False)\n",
    "\n",
    "    Hessian_diag = []\n",
    "    for param in mnist_model.parameters():\n",
    "        ps = param.size()\n",
    "        print(\"parameter size: \", ps)\n",
    "        Hessian_diag.append(torch.zeros(ps, device=device))\n",
    "        #print(param.numel())\n",
    "\n",
    "    tau = 1/var0\n",
    "    max_len = len(train_loader)\n",
    "\n",
    "    with backpack(DiagHessian()):\n",
    "\n",
    "        for batch_idx, (x, y) in enumerate(train_loader):\n",
    "\n",
    "            if device == 'cuda':\n",
    "                x, y = x.cuda(), y.cuda()\n",
    "\n",
    "            model.zero_grad()\n",
    "            loss = lossfunc(model(x), y)\n",
    "            loss.backward()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                # Hessian of weight\n",
    "                for idx, param in enumerate(model.parameters()):\n",
    "\n",
    "                    H_ = param.diag_h\n",
    "                    #add bias here\n",
    "                    H_ += tau * torch.ones(H_.size(), device=device)\n",
    "\n",
    "                    rho = min(1-1/(batch_idx+1), 0.995)\n",
    "\n",
    "                    Hessian_diag[idx] = rho*Hessian_diag[idx] + (1-rho)*H_\n",
    "            \n",
    "            if verbose:\n",
    "                print(\"Batch: {}/{}\".format(batch_idx, max_len))\n",
    "\n",
    "    #combine all elements of the Hessian to one big vector\n",
    "    Hessian_diag = torch.cat([el.view(-1) for el in Hessian_diag])\n",
    "    print(\"Hessian_size: \", Hessian_diag.size())\n",
    "    num_params = np.sum([p.numel() for p in model.parameters()])\n",
    "    assert(num_params == Hessian_diag.size(-1))\n",
    "    return(Hessian_diag)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parameter size:  torch.Size([32, 1, 5, 5])\n",
      "parameter size:  torch.Size([32])\n",
      "parameter size:  torch.Size([64, 32, 5, 5])\n",
      "parameter size:  torch.Size([64])\n",
      "parameter size:  torch.Size([256, 1024])\n",
      "parameter size:  torch.Size([256])\n",
      "parameter size:  torch.Size([10, 256])\n",
      "parameter size:  torch.Size([10])\n",
      "Hessian_size:  torch.Size([317066])\n"
     ]
    }
   ],
   "source": [
    "Hessian_MNIST = get_Hessian_NN(model=mnist_model, train_loader=mnist_train_loader, var0=200, verbose=False, device='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0056, 0.0059, 0.0067,  ..., 0.0068, 0.0074, 0.0070], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "print(Hessian_MNIST)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "## play around with Backpack\n",
    "def get_Hessian_KFAC_NN(model, train_loader, var0, device='cpu', verbose=True):\n",
    "    lossfunc = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    extend(lossfunc, debug=False)\n",
    "    extend(model, debug=False)\n",
    "\n",
    "    Hessian_KFAC_U = []\n",
    "    Hessian_KFAC_V = []\n",
    "    for param in mnist_model.parameters():\n",
    "        ps = param.size()\n",
    "        print(\"parameter size: \", ps)\n",
    "        Hessian_KFAC_U.append(torch.zeros(ps, device=device))\n",
    "        Hessian_KFAC_V.append(torch.zeros(ps, device=device))\n",
    "        #print(param.numel())\n",
    "\n",
    "    tau = 1/var0\n",
    "    max_len = len(train_loader)\n",
    "\n",
    "    with backpack(KFAC()):\n",
    "\n",
    "        for batch_idx, (x, y) in enumerate(train_loader):\n",
    "\n",
    "            if device == 'cuda':\n",
    "                x, y = x.cuda(), y.cuda()\n",
    "\n",
    "            model.zero_grad()\n",
    "            loss = lossfunc(model(x), y)\n",
    "            loss.backward()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                # Hessian of weight\n",
    "                for idx, param in enumerate(model.parameters()):\n",
    "\n",
    "                    # Hessian of weight\n",
    "                    U_, V_ = param.kfac\n",
    "\n",
    "                    U_ = np.sqrt(batch_size)*U_ + np.sqrt(tau)*torch.eye(m, device=device)\n",
    "                    V_ = np.sqrt(batch_size)*V_ + np.sqrt(tau)*torch.eye(n, device=device)\n",
    "\n",
    "                    rho = min(1-1/(batch_idx+1), 0.95)\n",
    "\n",
    "                    Hessian_KFAC_U[idx] = rho*Hessian_KFAC_U + (1-rho)*U_\n",
    "                    Hessian_KFAC_V[idx] = rho*Hessian_KFAC_V + (1-rho)*V_\n",
    "            \n",
    "            if verbose:\n",
    "                print(\"Batch: {}/{}\".format(batch_idx, max_len))\n",
    "\n",
    "    #combine all elements of the Hessian to one big vector\n",
    "    Hessian_diag = torch.cat([el.view(-1) for el in Hessian_diag])\n",
    "    print(\"Hessian_size: \", Hessian_diag.size())\n",
    "    num_params = np.sum([p.numel() for p in model.parameters()])\n",
    "    assert(num_params == Hessian_diag.size(-1))\n",
    "    return(Hessian_diag)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_jacobians_with_backpack(model, x, y, lossfunc):\n",
    "    \"\"\"\n",
    "    Returns the jacobians of the network\n",
    "\n",
    "    The output is a list. Each element in the list is a tensor\n",
    "    corresponding to the model.parameters().\n",
    "\n",
    "    The tensor are of the form [N, *, C] where N is the batch dimension,\n",
    "    C is the number of classes (output size of the network)\n",
    "    and * is the shape of the model parameters\n",
    "    \"\"\"\n",
    "    loss = lossfunc(model(x), y)\n",
    "    print(loss)\n",
    "\n",
    "    with backpack(NetJac()):\n",
    "        loss.backward()\n",
    "\n",
    "    jacs = []\n",
    "    for p in model.parameters():\n",
    "        jacs.append(p.netjacs.data.detach())\n",
    "    return jacs\n",
    "\n",
    "def transform2full_jac(backpack_jacobian):\n",
    "\n",
    "    jac_full = []\n",
    "    #batch_size\n",
    "    N = backpack_jacobian[0].size(0)\n",
    "    #num classes\n",
    "    k = backpack_jacobian[0].size(-1)\n",
    "    for j in backpack_jacobian:\n",
    "        jac_full.append(j.view(N, -1, k).permute(0,2,1))\n",
    "    jac_full = torch.cat(jac_full, dim=-1)\n",
    "    return(jac_full)\n",
    "\n",
    "def get_Jacobian(model, x, y, lossfunc):\n",
    "    return(transform2full_jac(compute_jacobians_with_backpack(model, x, y, lossfunc)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_Diagonal_full(model, test_loader, Hessian, verbose=True, num_samples=100, cuda=False, timing=False):\n",
    "    \n",
    "    lossfunc = torch.nn.CrossEntropyLoss()\n",
    "    extend(lossfunc, debug=False)\n",
    "    \n",
    "    py = []\n",
    "    if timing:\n",
    "        time_sum = 0\n",
    "    \n",
    "    max_len = len(test_loader)\n",
    "    for batch_idx, (x, y) in enumerate(test_loader):\n",
    "        \n",
    "        if cuda:\n",
    "            x, y = x.cuda(), y.cuda()\n",
    "        \n",
    "        J = get_Jacobian(model, x, y, lossfunc)\n",
    "        J = J.detach()\n",
    "        batch_size = J.size(0)\n",
    "        num_classes = J.size(1)\n",
    "        Cov_pred = torch.bmm(J * Hessian, J.permute(0, 2, 1))\n",
    "        Cov_pred = Cov_pred.detach()\n",
    "        if verbose:\n",
    "            print(\"Jacobian size: \", J.size())\n",
    "            print(\"cov pred size: \", Cov_pred.size())\n",
    "        \n",
    "        mu_pred = model(x).detach()\n",
    "        post_pred = MultivariateNormal(mu_pred, Cov_pred)\n",
    "\n",
    "        # MC-integral\n",
    "        t0 = time.time()\n",
    "        py_ = 0\n",
    "\n",
    "        for _ in range(num_samples):\n",
    "            f_s = post_pred.rsample()\n",
    "            py_ += torch.softmax(f_s, 1)\n",
    "\n",
    "\n",
    "        py_ /= num_samples\n",
    "        py_ = py_.detach()\n",
    "\n",
    "        py.append(py_)\n",
    "        t1 = time.time()\n",
    "        if timing:\n",
    "            time_sum += (t1-t0)\n",
    "\n",
    "        if verbose:\n",
    "            print(\"Batch: {}/{}\".format(batch_idx, max_len))\n",
    "    \n",
    "    if timing:\n",
    "        print(\"total time used for transform: {:.05f}\".format(time_sum))\n",
    "\n",
    "    return torch.cat(py, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE_TEST_FMNIST = 32#64#128\n",
    "BATCH_SIZE_TEST_KMNIST = 32#64#128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "FMNIST_test = torchvision.datasets.FashionMNIST(\n",
    "        '~/data/fmnist', train=False, download=True,\n",
    "        transform=MNIST_transform)   #torchvision.transforms.ToTensor())\n",
    "\n",
    "FMNIST_test_loader = torch.utils.data.DataLoader(\n",
    "    FMNIST_test,\n",
    "    batch_size=BATCH_SIZE_TEST_FMNIST, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "KMNIST_test = torchvision.datasets.KMNIST(\n",
    "        '~/data/kmnist', train=False, download=True,\n",
    "        transform=MNIST_transform)\n",
    "\n",
    "KMNIST_test_loader = torch.utils.data.DataLoader(\n",
    "    KMNIST_test,\n",
    "    batch_size=BATCH_SIZE_TEST_KMNIST, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Load notMNIST\"\"\"\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from matplotlib.pyplot import imread\n",
    "from torch import Tensor\n",
    "\n",
    "\"\"\"\n",
    "Loads the train/test set. \n",
    "Every image in the dataset is 28x28 pixels and the labels are numbered from 0-9\n",
    "for A-J respectively.\n",
    "Set root to point to the Train/Test folders.\n",
    "\"\"\"\n",
    "\n",
    "# Creating a sub class of torch.utils.data.dataset.Dataset\n",
    "class notMNIST(Dataset):\n",
    "\n",
    "    # The init method is called when this class will be instantiated\n",
    "    def __init__(self, root, transform):\n",
    "        \n",
    "        #super(notMNIST, self).__init__(root, transform=transform)\n",
    "\n",
    "        self.transform = transform\n",
    "        \n",
    "        Images, Y = [], []\n",
    "        folders = os.listdir(root)\n",
    "\n",
    "        for folder in folders:\n",
    "            folder_path = os.path.join(root, folder)\n",
    "            for ims in os.listdir(folder_path):\n",
    "                try:\n",
    "                    img_path = os.path.join(folder_path, ims)\n",
    "                    Images.append(np.array(imread(img_path)))\n",
    "                    Y.append(ord(folder) - 65)  # Folders are A-J so labels will be 0-9\n",
    "                except:\n",
    "                    # Some images in the dataset are damaged\n",
    "                    print(\"File {}/{} is broken\".format(folder, ims))\n",
    "        data = [(x, y) for x, y in zip(Images, Y)]\n",
    "        self.data = data\n",
    "        self.targets = torch.Tensor(Y)\n",
    "\n",
    "    # The number of items in the dataset\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    # The Dataloader is a generator that repeatedly calls the getitem method.\n",
    "    # getitem is supposed to return (X, Y) for the specified index.\n",
    "    def __getitem__(self, index):\n",
    "        img = self.data[index][0]\n",
    "\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "            \n",
    "        # Input for Conv2D should be Channels x Height x Width\n",
    "        img_tensor = Tensor(img).view(1, 28, 28).float()\n",
    "        label = self.data[index][1]\n",
    "        return (img_tensor, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File F/Q3Jvc3NvdmVyIEJvbGRPYmxpcXVlLnR0Zg==.png is broken\n",
      "File A/RGVtb2NyYXRpY2FCb2xkT2xkc3R5bGUgQm9sZC50dGY=.png is broken\n"
     ]
    }
   ],
   "source": [
    "#root = os.path.abspath('~/data')\n",
    "root = os.path.expanduser('~/data')\n",
    "\n",
    "# Instantiating the notMNIST dataset class we created\n",
    "notMNIST_test = notMNIST(root=os.path.join(root, 'notMNIST_small'),\n",
    "                               transform=MNIST_transform)\n",
    "\n",
    "# Creating a dataloader\n",
    "not_mnist_test_loader = torch.utils.data.dataloader.DataLoader(\n",
    "                            dataset=notMNIST_test,\n",
    "                            batch_size=BATCH_SIZE_TEST_KMNIST,\n",
    "                            shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MAP estimate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "targets = MNIST_test.targets.numpy()\n",
    "targets_FMNIST = FMNIST_test.targets.numpy()\n",
    "targets_notMNIST = notMNIST_test.targets.numpy().astype(int)\n",
    "targets_KMNIST = KMNIST_test.targets.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_test_in_MAP = predict_MAP(mnist_model, mnist_test_loader, cuda=True).cpu().numpy()\n",
    "mnist_test_out_fmnist_MAP = predict_MAP(mnist_model, FMNIST_test_loader, cuda=True).cpu().numpy()\n",
    "mnist_test_out_notMNIST_MAP = predict_MAP(mnist_model, not_mnist_test_loader, cuda=True).cpu().numpy()\n",
    "mnist_test_out_KMNIST_MAP = predict_MAP(mnist_model, KMNIST_test_loader, cuda=True).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_in_MAP, prob_correct_in_MAP, ent_in_MAP, MMC_in_MAP = get_in_dist_values(mnist_test_in_MAP, targets)\n",
    "acc_out_FMNIST_MAP, prob_correct_out_FMNIST_MAP, ent_out_FMNIST_MAP, MMC_out_FMNIST_MAP, auroc_out_FMNIST_MAP = get_out_dist_values(mnist_test_in_MAP, mnist_test_out_fmnist_MAP, targets_FMNIST)\n",
    "acc_out_notMNIST_MAP, prob_correct_out_notMNIST_MAP, ent_out_notMNIST_MAP, MMC_out_notMNIST_MAP, auroc_out_notMNIST_MAP = get_out_dist_values(mnist_test_in_MAP, mnist_test_out_notMNIST_MAP, targets_notMNIST)\n",
    "acc_out_KMNIST_MAP, prob_correct_out_KMNIST_MAP, ent_out_KMNIST_MAP, MMC_out_KMNIST_MAP, auroc_out_KMNIST_MAP = get_out_dist_values(mnist_test_in_MAP, mnist_test_out_KMNIST_MAP, targets_KMNIST)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[In, MAP, mnist] Accuracy: 0.990; average entropy: 0.035;     MMC: 0.989; Prob @ correct: 0.100\n",
      "[Out-MAP, KFAC, FMNIST] Accuracy: 0.108; Average entropy: 1.390;    MMC: 0.514; AUROC: 0.993; Prob @ correct: 0.100\n",
      "[Out-MAP, KFAC, notMNIST] Accuracy: 0.108; Average entropy: 0.808;    MMC: 0.709; AUROC: 0.956; Prob @ correct: 0.100\n",
      "[Out-MAP, KFAC, KMNIST] Accuracy: 0.083; Average entropy: 0.861;    MMC: 0.688; AUROC: 0.973; Prob @ correct: 0.100\n"
     ]
    }
   ],
   "source": [
    "print_in_dist_values(acc_in_MAP, prob_correct_in_MAP, ent_in_MAP, MMC_in_MAP, 'mnist', 'MAP')\n",
    "print_out_dist_values(acc_out_FMNIST_MAP, prob_correct_out_FMNIST_MAP, ent_out_FMNIST_MAP, MMC_out_FMNIST_MAP, auroc_out_FMNIST_MAP, 'FMNIST', 'MAP')\n",
    "print_out_dist_values(acc_out_notMNIST_MAP, prob_correct_out_notMNIST_MAP, ent_out_notMNIST_MAP, MMC_out_notMNIST_MAP, auroc_out_notMNIST_MAP, 'notMNIST', 'MAP')\n",
    "print_out_dist_values(acc_out_KMNIST_MAP, prob_correct_out_KMNIST_MAP, ent_out_KMNIST_MAP, MMC_out_KMNIST_MAP, auroc_out_KMNIST_MAP, 'KMNIST', 'MAP')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy: 0.990 with std 0.000\n",
      "MMC in: 0.989 with std 0.000\n",
      "MMC out fmnist: 0.514 with std 0.000\n",
      "MMC out notmnist: 0.709 with std 0.000\n",
      "MMC out kmnist: 0.688 with std 0.000\n",
      "AUROC out fmnist: 0.993 with std 0.000\n",
      "AUROC out notmnist: 0.956 with std 0.000\n",
      "AUROC out kmnist: 0.973 with std 0.000\n"
     ]
    }
   ],
   "source": [
    "#MAP estimate\n",
    "#seeds are 123,124,125,126,127\n",
    "acc_in = [0.990]\n",
    "mmc_in = [0.989]\n",
    "mmc_out_fmnist = [0.514]\n",
    "mmc_out_notmnist = [0.709]\n",
    "mmc_out_kmnist = [0.688]\n",
    "\n",
    "auroc_out_fmnist = [0.993]\n",
    "auroc_out_notmnist = [0.956]\n",
    "auroc_out_kmnist = [0.973]\n",
    "\n",
    "print(\"accuracy: {:.03f} with std {:.03f}\".format(np.mean(acc_in), np.std(acc_in)))\n",
    "\n",
    "print(\"MMC in: {:.03f} with std {:.03f}\".format(np.mean(mmc_in), np.std(mmc_in)))\n",
    "print(\"MMC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_fmnist), np.std(mmc_out_fmnist)))\n",
    "print(\"MMC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_notmnist), np.std(mmc_out_notmnist)))\n",
    "print(\"MMC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_kmnist), np.std(mmc_out_kmnist)))\n",
    "\n",
    "print(\"AUROC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_fmnist), np.std(auroc_out_fmnist)))\n",
    "print(\"AUROC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_notmnist), np.std(auroc_out_notmnist)))\n",
    "print(\"AUROC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_kmnist), np.std(auroc_out_kmnist)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diag Hessian Sampling estimate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0100, device='cuda:0', grad_fn=<NllLossBackward>)\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "('Compared shapes [10, 10] and [32, 10] do not match. ', 'Got [32, 10, 10] and [32, 10]')",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-41-520708a9d980>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmnist_test_in_D\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredict_Diagonal_full\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmnist_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmnist_test_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mHessian_MNIST\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcuda\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtiming\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mmnist_test_out_FMNIST_D\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredict_Diagonal_full\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmnist_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mFMNIST_test_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mHessian_MNIST\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcuda\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtiming\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mmnist_test_out_notMNIST_D\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredict_Diagonal_full\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmnist_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnot_mnist_test_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mHessian_MNIST\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcuda\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtiming\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mmnist_test_out_KMNIST_D\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredict_Diagonal_full\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmnist_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mKMNIST_test_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mHessian_MNIST\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcuda\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtiming\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-29-e0d7052531b0>\u001b[0m in \u001b[0;36mpredict_Diagonal_full\u001b[0;34m(model, test_loader, Hessian, verbose, num_samples, cuda, timing)\u001b[0m\n\u001b[1;32m     14\u001b[0m             \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m         \u001b[0mJ\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_Jacobian\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlossfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m         \u001b[0mJ\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mJ\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0mbatch_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mJ\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-28-2e244c663457>\u001b[0m in \u001b[0;36mget_Jacobian\u001b[0;34m(model, x, y, lossfunc)\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_Jacobian\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlossfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m     \u001b[0;32mreturn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtransform2full_jac\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcompute_jacobians_with_backpack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlossfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-28-2e244c663457>\u001b[0m in \u001b[0;36mcompute_jacobians_with_backpack\u001b[0;34m(model, x, y, lossfunc)\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mbackpack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNetJac\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m     \u001b[0mjacs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    196\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    197\u001b[0m         \"\"\"\n\u001b[0;32m--> 198\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    199\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    200\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     98\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     99\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 100\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    102\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: ('Compared shapes [10, 10] and [32, 10] do not match. ', 'Got [32, 10, 10] and [32, 10]')"
     ]
    }
   ],
   "source": [
    "mnist_test_in_D = predict_Diagonal_full(mnist_model, mnist_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True, num_samples=num_samples).cpu().numpy()\n",
    "mnist_test_out_FMNIST_D = predict_Diagonal_full(mnist_model, FMNIST_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True, num_samples=num_samples).cpu().numpy()\n",
    "mnist_test_out_notMNIST_D = predict_Diagonal_full(mnist_model, not_mnist_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True, num_samples=num_samples).cpu().numpy()\n",
    "mnist_test_out_KMNIST_D = predict_Diagonal_full(mnist_model, KMNIST_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True, num_samples=num_samples).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_in_D, prob_correct_in_D, ent_in_D, MMC_in_D = get_in_dist_values(mnist_test_in_D, targets)\n",
    "acc_out_FMNIST_D, prob_correct_out_FMNIST_D, ent_out_FMNIST_D, MMC_out_FMNIST_D, auroc_out_FMNIST_D = get_out_dist_values(mnist_test_in_D, mnist_test_out_FMNIST_D, targets_FMNIST)\n",
    "acc_out_notMNIST_D, prob_correct_out_notMNIST_D, ent_out_notMNIST_D, MMC_out_notMNIST_D, auroc_out_notMNIST_D = get_out_dist_values(mnist_test_in_D, mnist_test_out_notMNIST_D, targets_notMNIST)\n",
    "acc_out_KMNIST_D, prob_correct_out_KMNIST_D, ent_out_KMNIST_D, MMC_out_KMNIST_D, auroc_out_KMNIST_D = get_out_dist_values(mnist_test_in_D, mnist_test_out_KMNIST_D, targets_KMNIST)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_in_dist_values(acc_in_D, prob_correct_in_D, ent_in_D, MMC_in_D, 'mnist', 'Diag')\n",
    "print_out_dist_values(acc_out_FMNIST_D, prob_correct_out_FMNIST_D, ent_out_FMNIST_D, MMC_out_FMNIST_D, auroc_out_FMNIST_D, test='fmnist', method='Diag')\n",
    "print_out_dist_values(acc_out_notMNIST_D, prob_correct_out_notMNIST_D, ent_out_notMNIST_D, MMC_out_notMNIST_D, auroc_out_notMNIST_D, test='notMNIST', method='Diag')\n",
    "print_out_dist_values(acc_out_KMNIST_D, prob_correct_out_KMNIST_D, ent_out_KMNIST_D, MMC_out_KMNIST_D, auroc_out_KMNIST_D, test='KMNIST', method='Diag')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Diag Sampling\n",
    "#seeds are 123,124,125,126,127\n",
    "time_lpb_in = [26.45687, 27.79014, 26.57382, 26.47316, 26.88896]\n",
    "time_lpb_out_fmnist = [26.78206, 27.77104, 26.32712, 26.37834, 26.62998]\n",
    "time_lpb_out_notmnist = [50.15894, 52.16657, 49.43548, 49.79360, 49.76200]\n",
    "time_lpb_out_kmnist = [26.71048, 27.93303, 26.52897, 26.47667, 26.88779]\n",
    "\n",
    "acc_in = [0.990, 0.990, 0.990, 0.990, 0.990]\n",
    "mmc_in = [0.942, 0.941, 0.942, 0.942, 0.941]\n",
    "mmc_out_fmnist = [0.398, 0.397, 0.397, 0.398, 0.396]\n",
    "mmc_out_notmnist = [0.543, 0.542, 0.543, 0.543, 0.542]\n",
    "mmc_out_kmnist = [0.514, 0.512, 0.513, 0.514, 0.512]\n",
    "\n",
    "auroc_out_fmnist = [0.992, 0.992, 0.992, 0.992, 0.992]\n",
    "auroc_out_notmnist = [0.960, 0.960, 0.960, 0.960, 0.960]\n",
    "auroc_out_kmnist = [0.974, 0.974, 0.974, 0.974, 0.974]\n",
    "\n",
    "print(\"Sampling Bridge time in: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_in), np.std(time_lpb_in)))\n",
    "print(\"Sampling Bridge time out fmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_fmnist), np.std(time_lpb_out_fmnist)))\n",
    "print(\"Sampling Bridge time out notmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_notmnist), np.std(time_lpb_out_notmnist)))\n",
    "print(\"Sampling Bridge time out kmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_kmnist), np.std(time_lpb_out_kmnist)))\n",
    "\n",
    "print(\"accuracy: {:.03f} with std {:.03f}\".format(np.mean(acc_in), np.std(acc_in)))\n",
    "\n",
    "print(\"MMC in: {:.03f} with std {:.03f}\".format(np.mean(mmc_in), np.std(mmc_in)))\n",
    "print(\"MMC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_fmnist), np.std(mmc_out_fmnist)))\n",
    "print(\"MMC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_notmnist), np.std(mmc_out_notmnist)))\n",
    "print(\"MMC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_kmnist), np.std(mmc_out_kmnist)))\n",
    "\n",
    "print(\"AUROC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_fmnist), np.std(auroc_out_fmnist)))\n",
    "print(\"AUROC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_notmnist), np.std(auroc_out_notmnist)))\n",
    "print(\"AUROC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_kmnist), np.std(auroc_out_kmnist)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full KFAC Laplace Approximation sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The KFAC Hessian (and even its components) is too large to store in memory and therefore skipped for this condition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dirichlet Laplace Approximation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_test_in_DIR_LPA = predict_DIR_LPA(mnist_model, mnist_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True).cpu().numpy()\n",
    "mnist_test_out_FMNIST_DIR_LPA = predict_DIR_LPA(mnist_model, FMNIST_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True).cpu().numpy()\n",
    "mnist_test_out_notMNIST_DIR_LPA = predict_DIR_LPA(mnist_model, not_mnist_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True).cpu().numpy()\n",
    "mnist_test_out_KMNIST_DIR_LPA = predict_DIR_LPA(mnist_model, KMNIST_test_loader, Hessian_MNIST, verbose=False, cuda=True, timing=True).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_test_in_DIR_LPAn = mnist_test_in_DIR_LPA/mnist_test_in_DIR_LPA.sum(1).reshape(-1,1)\n",
    "mnist_test_out_FMNIST_DIR_LPAn = mnist_test_out_FMNIST_DIR_LPA/mnist_test_out_FMNIST_DIR_LPA.sum(1).reshape(-1,1)\n",
    "mnist_test_out_notMNIST_DIR_LPAn = mnist_test_out_notMNIST_DIR_LPA/mnist_test_out_notMNIST_DIR_LPA.sum(1).reshape(-1,1)\n",
    "mnist_test_out_KMNIST_DIR_LPAn = mnist_test_out_KMNIST_DIR_LPA/mnist_test_out_KMNIST_DIR_LPA.sum(1).reshape(-1,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_in_DIR_LPAn, prob_correct_in_DIR_LPAn, ent_in_DIR_LPAn, MMC_in_DIR_LPAn = get_in_dist_values(mnist_test_in_DIR_LPAn, targets)\n",
    "acc_out_FMNIST_DIR_LPAn, prob_correct_out_FMNIST_DIR_LPAn, ent_out_FMNIST_DIR_LPAn, MMC_out_FMNIST_DIR_LPAn, auroc_out_FMNIST_DIR_LPAn = get_out_dist_values(mnist_test_in_DIR_LPAn, mnist_test_out_FMNIST_DIR_LPAn, targets_FMNIST)\n",
    "acc_out_notMNIST_DIR_LPAn, prob_correct_out_notMNIST_DIR_LPAn, ent_out_notMNIST_DIR_LPAn, MMC_out_notMNIST_DIR_LPAn, auroc_out_notMNIST_DIR_LPAn = get_out_dist_values(mnist_test_in_DIR_LPAn, mnist_test_out_notMNIST_DIR_LPAn, targets_notMNIST)\n",
    "acc_out_KMNIST_DIR_LPAn, prob_correct_out_KMNIST_DIR_LPAn, ent_out_KMNIST_DIR_LPAn, MMC_out_KMNIST_DIR_LPAn, auroc_out_KMNIST_DIR_LPAn = get_out_dist_values(mnist_test_in_DIR_LPAn, mnist_test_out_KMNIST_DIR_LPAn, targets_KMNIST)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_in_dist_values(acc_in_DIR_LPAn, prob_correct_in_DIR_LPAn, ent_in_DIR_LPAn, MMC_in_DIR_LPAn, 'mnist', 'DIR_LPAn')\n",
    "print_out_dist_values(acc_out_FMNIST_DIR_LPAn, prob_correct_out_FMNIST_DIR_LPAn, ent_out_FMNIST_DIR_LPAn, MMC_out_FMNIST_DIR_LPAn, auroc_out_FMNIST_DIR_LPAn, test='fmnist', method='DIR_LPAn')\n",
    "print_out_dist_values(acc_out_notMNIST_DIR_LPAn, prob_correct_out_notMNIST_DIR_LPAn, ent_out_notMNIST_DIR_LPAn, MMC_out_notMNIST_DIR_LPAn, auroc_out_notMNIST_DIR_LPAn, test='notMNIST', method='DIR_LPAn')\n",
    "print_out_dist_values(acc_out_KMNIST_DIR_LPAn, prob_correct_out_KMNIST_DIR_LPAn, ent_out_KMNIST_DIR_LPAn, MMC_out_KMNIST_DIR_LPAn, auroc_out_KMNIST_DIR_LPAn, test='KMNIST', method='DIR_LPAn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Laplace Bridge\n",
    "#seeds are 123,124,125,126,127\n",
    "time_lpb_in = [0.06198, 0.06407, 0.06117, 0.06193, 0.06098]\n",
    "time_lpb_out_fmnist = [0.06195, 0.06337, 0.06156, 0.06214, 0.06095]\n",
    "time_lpb_out_notmnist = [0.11700, 0.12246, 0.11510, 0.11567, 0.11433]\n",
    "time_lpb_out_kmnist = [0.06194, 0.06504, 0.06136, 0.06200, 0.06136]\n",
    "\n",
    "acc_in = [0.990, 0.990, 0.990, 0.990, 0.990]\n",
    "mmc_in = [0.987, 0.987, 0.987, 0.987, 0.987]\n",
    "mmc_out_fmnist = [0.363, 0.363, 0.363, 0.363, 0.363]\n",
    "mmc_out_notmnist = [0.649, 0.649, 0.649, 0.649, 0.649]\n",
    "mmc_out_kmnist = [0.637, 0.637, 0.637, 0.638, 0.637]\n",
    "\n",
    "auroc_out_fmnist = [0.996, 0.996, 0.996, 0.996, 0.996]\n",
    "auroc_out_notmnist = [0.961, 0.961, 0.961, 0.961, 0.961]\n",
    "auroc_out_kmnist = [0.973, 0.973, 0.973, 0.973, 0.973]\n",
    "\n",
    "print(\"Laplace Bridge time in: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_in), np.std(time_lpb_in)))\n",
    "print(\"Laplace Bridge time out fmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_fmnist), np.std(time_lpb_out_fmnist)))\n",
    "print(\"Laplace Bridge time out notmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_notmnist), np.std(time_lpb_out_notmnist)))\n",
    "print(\"Laplace Bridge time out kmnist: {:.03f} with std {:.03f}\".format(np.mean(time_lpb_out_kmnist), np.std(time_lpb_out_kmnist)))\n",
    "\n",
    "print(\"accuracy: {:.03f} with std {:.03f}\".format(np.mean(acc_in), np.std(acc_in)))\n",
    "\n",
    "print(\"MMC in: {:.03f} with std {:.03f}\".format(np.mean(mmc_in), np.std(mmc_in)))\n",
    "print(\"MMC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_fmnist), np.std(mmc_out_fmnist)))\n",
    "print(\"MMC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_notmnist), np.std(mmc_out_notmnist)))\n",
    "print(\"MMC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(mmc_out_kmnist), np.std(mmc_out_kmnist)))\n",
    "\n",
    "print(\"AUROC out fmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_fmnist), np.std(auroc_out_fmnist)))\n",
    "print(\"AUROC out notmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_notmnist), np.std(auroc_out_notmnist)))\n",
    "print(\"AUROC out kmnist: {:.03f} with std {:.03f}\".format(np.mean(auroc_out_kmnist), np.std(auroc_out_kmnist)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if condition holds\n",
    "\n",
    "def check_condition(alpha_vecs):\n",
    "    #note that this is vectorized\n",
    "    alpha_sum = alpha_vecs.sum(1)\n",
    "    alpha_max = alpha_vecs.max(1)\n",
    "    alpha_sum_minus = alpha_sum - alpha_max\n",
    "    right_side = 0.25 * (np.sqrt(9 * alpha_sum_minus**2 + 10 * alpha_sum_minus + 1) - alpha_sum_minus - 1)\n",
    "    cases = alpha_max > right_side\n",
    "    percentage = np.sum(cases)/len(cases)\n",
    "    return(percentage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.sum(check_condition(mnist_test_in_DIR_LPA)))\n",
    "print(np.sum(check_condition(mnist_test_out_FMNIST_DIR_LPA)))\n",
    "print(np.sum(check_condition(mnist_test_out_notMNIST_DIR_LPA)))\n",
    "print(np.sum(check_condition(mnist_test_out_KMNIST_DIR_LPA)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# since our package is not compatible with pytorch anymore the results have been conducted with a last-layer\n",
    "# Laplace approximation of the network in the other jupyter notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
